#!/usr/bin/env python
from __future__ import print_function

import tensorflow as tf
import cv2
import sys
sys.path.append("game/")
import wrapped_flappy_bird as game
import random
import numpy as np
from collections import deque
IMGWIDTH = 80
IMGHEIGHT =  80
GAME = 'bird' # the name of the game being played for log files
ACTIONS = 2 # number of valid actions
GAMMA = 0.99 # decay rate of past observations
OBSERVE = 100. # timesteps to observe before training
EXPLORE = 20000. # frames over which to anneal epsilon
FINAL_EPSILON = 0.0001 # final value of epsilon
INITIAL_EPSILON = 0.01 # starting value of epsilon
REPLAY_MEMORY = 10000 # number of previous transitions to remember
BATCH = 32 # size of minibatch
FRAME_PER_ACTION = 1
TARGETALPHA = 0.5
UPDATE_FRAME = 8
UPDATE_TARGET_FRAME = 24

def weight_variable(shape,name=None):
    initial = tf.truncated_normal(shape, stddev = 0.01)
    return tf.Variable(initial,name)

def bias_variable(shape,name=None):
    initial = tf.constant(0.01, shape = shape)
    return tf.Variable(initial,name)

def conv2d(x, W, stride):
    return tf.nn.conv2d(x, W, strides = [1, stride, stride, 1], padding = "SAME")

def max_pool_2x2(x):
    return tf.nn.max_pool(x, ksize = [1, 2, 2, 1], strides = [1, 2, 2, 1], padding = "SAME")

def createNetwork(i):
    # network weights
    W_conv1 = weight_variable([8, 8, 4, 32])
    b_conv1 = bias_variable([32])

    W_conv2 = weight_variable([4, 4, 32, 64])
    b_conv2 = bias_variable([64])

    W_conv3 = weight_variable([3, 3, 64, 64])
    b_conv3 = bias_variable([64])

    n = str(i);
    Wfc1 = weight_variable([1600, 512],name='wfc1'+n)
    bfc1 = bias_variable([512],name='bfc1'+n)

    # W_fc2 = weight_variable([512, ACTIONS])
    # b_fc2 = bias_variable([ACTIONS])
    AW1 = weight_variable([256,ACTIONS],name='aw1'+n)
    VW1 = weight_variable([256,1],name='vw1'+n)

    # input layer
    s = tf.placeholder("float", [None, IMGWIDTH, IMGHEIGHT, 4])

    # hidden layers
    # (?, 20, 20, 32)
    h_conv1 = tf.nn.relu(conv2d(s, W_conv1, 4) + b_conv1)
    # (?, 10, 10, 32)
    h_pool1 = max_pool_2x2(h_conv1)
    # (?, 5, 5, 64)
    h_conv2 = tf.nn.relu(conv2d(h_pool1, W_conv2, 2) + b_conv2)
    #h_pool2 = max_pool_2x2(h_conv2)
    h_conv3 = tf.nn.relu(conv2d(h_conv2, W_conv3, 1) + b_conv3)
    #h_pool3 = max_pool_2x2(h_conv3)

    h_conv3_flat = tf.reshape(h_conv3, [-1, 1600])

    h_fc1 = tf.nn.relu(tf.matmul(h_conv3_flat, Wfc1) + bfc1)
    #[None,256]
    AC,VC = tf.split(1,2,h_fc1)
    # [None,ACTIONS]
    Adv = tf.matmul(AC,AW1)
    # [None,1]
    Value = tf.matmul(VC,VW1)
    # readout layer [None,Actions] q(s,a)
    #readout = tf.matmul(h_fc1, W_fc2) + b_fc2
    Q = Value +tf.subtract(Adv,tf.reduce_mean(Adv,reduction_indices=1,keep_dims=True))
    # readout = Adv
    return s, Q

def step(game_state,a):
    x, r, terminal = game_state.frame_step(a)
    x = cv2.cvtColor(cv2.resize(x, (IMGWIDTH, IMGHEIGHT)), cv2.COLOR_BGR2GRAY)
    ret, x = cv2.threshold(x, 1, 255, cv2.THRESH_BINARY)
    if terminal:
        r = -1
    return x,r,terminal

def trainNetwork(mainS, mainQ, targetS, targetQ, sess):
    # define the cost function
    a = tf.placeholder("float", [None, ACTIONS])
    y = tf.placeholder("float", [None])
    y_ = tf.reduce_sum(tf.multiply(mainQ, a), reduction_indices=1)
    cost = tf.reduce_mean(tf.square(y - y_))
    train_step = tf.train.AdamOptimizer(1e-6).minimize(cost)
    sess.run(tf.global_variables_initializer())
    ops = updateTargetGraph(tf.trainable_variables())
    saver,global_step = loadData(sess)
    # open up a game state to communicate with emulator
    game_state = game.GameState()

    # store the previous observations in replay memory
    D = deque()

    # printing
    # a_file = open("logs_" + GAME + "/readout.txt", 'w')
    # h_file = open("logs_" + GAME + "/hidden.txt", 'w')

    # get the first state by doing nothing and preprocess the image to 80x80x4
    do_nothing = np.zeros(ACTIONS)
    do_nothing[0] = 1
    terminalStep = 1
    totalStep = 0
    terminalTimes = 0
    x_t,r_0,terminal = step(game_state,do_nothing)
    s_t = np.stack((x_t, x_t, x_t, x_t), axis=2)
    # saving and loading networks


    # start training
    epsilon = INITIAL_EPSILON
    t = 0
    while True:
        # choose an action epsilon greedily
        readout_t = mainQ.eval(feed_dict={mainS : [s_t]})[0]
        a_t = np.zeros([ACTIONS])
        action_index = 0
        if t % FRAME_PER_ACTION == 0:
            if random.random() <= epsilon:
                action_index = random.randrange(ACTIONS)
                a_t[action_index] = 1
            else:
                action_index = np.argmax(readout_t)
                a_t[action_index] = 1
        else:
            a_t[0] = 1 # do nothing

        # scale down epsilon
        if epsilon > FINAL_EPSILON and t > OBSERVE:
            epsilon -= (INITIAL_EPSILON - FINAL_EPSILON) / EXPLORE

        # run the selected action and observe next state and reward
        x_t1,r_t,terminal = step(game_state,a_t)
        if terminal:
            totalStep += terminalStep
            terminalTimes += 1
            print('TerminalStep %d,Times:%d,AvgStep:%f' % (terminalStep,terminalTimes,totalStep/terminalTimes))
            terminalStep = 0
        else:
            terminalStep += 1
        x_t1 = np.reshape(x_t1, (IMGWIDTH, IMGHEIGHT, 1))
        #s_t1 = np.append(x_t1, s_t[:,:,1:], axis = 2)
        s_t1 = np.append(x_t1, s_t[:, :, :3], axis=2)

        # store the transition in D
        D.append((s_t, a_t, r_t, s_t1, terminal))
        if len(D) > REPLAY_MEMORY:
            D.popleft()

        # only train if done observing
        if t > OBSERVE and t % UPDATE_FRAME == 0:
            # sample a minibatch to train on
            minibatch = random.sample(D, BATCH)

            # get the batch variables
            s_j_batch = [d[0] for d in minibatch]
            a_batch = [d[1] for d in minibatch]
            r_batch = [d[2] for d in minibatch]
            s_j1_batch = [d[3] for d in minibatch]

            y_batch = []
            # readout_j1_batch = targetQ.eval(feed_dict = {targetS : s_j1_batch})
            readout_j1_batch = mainQ.eval(feed_dict={mainS: s_j1_batch})
            readout_j1_batch_target = targetQ.eval(feed_dict={targetS: s_j1_batch})
            for i in range(0, len(minibatch)):
                terminal = minibatch[i][4]
                # if terminal, only equals reward
                if terminal:
                    y_batch.append(r_batch[i])
                else:
                    maxIndex = np.argmax(readout_j1_batch[i])
                    y_batch.append(r_batch[i] + GAMMA * readout_j1_batch_target[i][maxIndex])
                    # y_batch.append(r_batch[i] + GAMMA * np.max(readout_j1_batch[i]))

            # perform gradient step
            train_step.run(feed_dict = {
                y : y_batch,
                a : a_batch,
                mainS : s_j_batch}
            )

            if t % UPDATE_TARGET_FRAME  == 0:
                print('maxtarget reward:%f' % np.max(y_batch))
                updateTarget(ops,sess)

        # update the old values
        s_t = s_t1
        t += 1

        # save progress every 10000 iterations
        if t % 1000 == 0:
            saver.save(sess, 'saved_networks/' + GAME + '-dqn', global_step = t + global_step)
        if t % 100 == 0:
            # print info
            state = ""
            if t <= OBSERVE:
                state = "observe"
            elif t > OBSERVE and t <= OBSERVE + EXPLORE:
                state = "explore"
            else:
                state = "train"

            print("TIMESTEP", t, "/ STATE", state, \
                "/ EPSILON", epsilon, "/ ACTION", action_index, "/ REWARD", r_t, \
                "/ Q_MAX %e" % np.max(readout_t))
        # write info to files
        '''
        if t % 10000 <= 100:
            a_file.write(",".join([str(x) for x in readout_t]) + '\n')
            h_file.write(",".join([str(x) for x in h_fc1.eval(feed_dict={s:[s_t]})[0]]) + '\n')
            cv2.imwrite("logs_tetris/frame" + str(t) + ".png", x_t1)
        '''

def updateTargetGraph(tfVar):
    num = len(tfVar)//2
    ops = []
    for idx,var in enumerate(tfVar[0:num]):
        ops.append(tfVar[idx + num].assign(var.value() * TARGETALPHA +(1 - TARGETALPHA) * tfVar[idx + num].value()))
    return ops

def updateTarget(ops,sess):
    for op in ops:
        sess.run(op)

def loadData(sess):
    saver = tf.train.Saver()
    checkpoint = tf.train.get_checkpoint_state("saved_networks")
    global_step = 0
    if checkpoint and checkpoint.model_checkpoint_path:
        saver.restore(sess, checkpoint.model_checkpoint_path)
        global_step = int(checkpoint.model_checkpoint_path.split('/')[-1].split('-')[-1])
        print("Successfully loaded:%s,global_step:%d"%(checkpoint.model_checkpoint_path,global_step))
    else:
        print("Could not find old network weights")
    return saver,global_step

def playGame():
    sess = tf.InteractiveSession()
    mainS, mainQ = createNetwork(0)
    targetS, targetQ = createNetwork(1)
    trainNetwork(mainS, mainQ,targetS,targetQ, sess)

def main():
    playGame()

if __name__ == "__main__":
    main()
